import os
import time
import glob
import random
import threading
import shutil
from pic_generator import process_images_with_boots

class TimeoutError(Exception):
    """超时异常"""
    pass

def timeout_handler(timeout_seconds):
    """超时处理器"""
    def decorator(func):
        def wrapper(*args, **kwargs):
            result = [None]
            exception = [None]
            
            def target():
                try:
                    result[0] = func(*args, **kwargs)
                except Exception as e:
                    exception[0] = e
            
            thread = threading.Thread(target=target)
            thread.daemon = True
            thread.start()
            thread.join(timeout_seconds)
            
            if thread.is_alive():
                raise TimeoutError(f"处理超时（超过{timeout_seconds}秒）")
            
            if exception[0]:
                raise exception[0]
            
            return result[0]
        return wrapper
    return decorator

def move_to_done_folder(original_image_path, generated_image_path):
    """
    将原图片和生成的图片移动到done文件夹
    
    Args:
        original_image_path (str): 原图片路径
        generated_image_path (str): 生成的图片路径
    """
    # 创建done文件夹（如果不存在）
    done_folder = "done"
    if not os.path.exists(done_folder):
        os.makedirs(done_folder)
        print(f"创建done文件夹: {done_folder}")
    
    # 移动原图片
    original_filename = os.path.basename(original_image_path)
    original_dest = os.path.join(done_folder, original_filename)
    shutil.move(original_image_path, original_dest)
    
    # 移动生成的图片
    generated_filename = os.path.basename(generated_image_path)
    generated_dest = os.path.join(done_folder, generated_filename)
    shutil.move(generated_image_path, generated_dest)

def move_to_other_folder(original_image_path):
    """
    将原图片移动到other文件夹（处理失败时）
    
    Args:
        original_image_path (str): 原图片路径
    """
    # 创建other文件夹（如果不存在）
    other_folder = "other"
    if not os.path.exists(other_folder):
        os.makedirs(other_folder)
        print(f"创建other文件夹: {other_folder}")
    
    # 移动原图片
    original_filename = os.path.basename(original_image_path)
    original_dest = os.path.join(other_folder, original_filename)
    shutil.move(original_image_path, original_dest)
    print(f"  ✓ 已移动到other文件夹: {original_filename}")

def batch_process_images_with_boots(boots_image_path, prompt=None, max_retries=3, timeout_seconds=300):
    """
    批量处理pic文件夹下的所有图片，为每张图片生成穿靴子的版本
    
    Args:
        boots_image_path (str): 靴子图片路径
        prompt (str, optional): 自定义提示词，如果不提供则使用默认提示词
        max_retries (int): 最大重试次数，默认3次
        timeout_seconds (int): 单次处理超时时间（秒），默认300秒（5分钟）
    
    Returns:
        dict: 处理结果统计
    """
    # 获取pic文件夹下所有支持的图片格式
    pic_folder = "pic"
    supported_formats = ["*.png", "*.jpg", "*.jpeg", "*.webp", "*.gif"]
    image_files = []
    
    for format_pattern in supported_formats:
        image_pattern = os.path.join(pic_folder, format_pattern)
        image_files.extend(glob.glob(image_pattern))
    
    if not image_files:
        print(f"在 {pic_folder} 文件夹中没有找到支持的图片格式")
        return {"success": 0, "failed": 0, "skipped": 0, "details": []}
    
    print(f"找到 {len(image_files)} 张图片，开始批量处理...")
    
    # 统计信息
    stats = {
        "success": 0,
        "failed": 0,
        "skipped": 0,
        "details": []
    }
    
    for i, image_path in enumerate(image_files, 1):
        print(f"\n[{i}/{len(image_files)}] 处理图片: {os.path.basename(image_path)}")
        
        # 检查是否已经存在生成的文件
        base_name = os.path.splitext(os.path.basename(image_path))[0]
        expected_output = os.path.join(pic_folder, f"{base_name}_ai.png")
        
        if os.path.exists(expected_output):
            print(f"  跳过: 文件 {os.path.basename(expected_output)} 已存在")
            stats["skipped"] += 1
            stats["details"].append({
                "input": image_path,
                "output": expected_output,
                "status": "skipped",
                "message": "文件已存在"
            })
            continue
        
        # 重试机制
        success = False
        last_error = None
        
        for attempt in range(max_retries):
            try:
                print(f"  尝试 {attempt + 1}/{max_retries}...")
                
                # 使用超时装饰器包装处理函数
                timeout_wrapper = timeout_handler(timeout_seconds)(process_images_with_boots)
                
                # 调用处理函数
                result = timeout_wrapper(image_path, boots_image_path, prompt)
                
                if result:
                    print(f"  ✓ 成功生成: {os.path.basename(result)}")
                    
                    # 移动原图片和生成的图片到done文件夹
                    try:
                        move_to_done_folder(image_path, result)
                        print(f"  ✓ 已移动到done文件夹")
                    except Exception as move_error:
                        print(f"  ⚠ 移动文件失败: {move_error}")
                    
                    stats["success"] += 1
                    stats["details"].append({
                        "input": image_path,
                        "output": result,
                        "status": "success",
                        "attempts": attempt + 1
                    })
                    success = True
                    break
                else:
                    raise Exception("process_images_with_boots 返回了 None")
                    
            except TimeoutError as e:
                last_error = str(e)
                print(f"  ✗ 尝试 {attempt + 1} 失败: {last_error}")
            except Exception as e:
                last_error = str(e)
                print(f"  ✗ 尝试 {attempt + 1} 失败: {last_error}")
                
                # 检查是否是空响应错误（直接跳过重试）
                if "received empty response from Gemini" in str(e) or "no meaningful content in candidates" in str(e) or "channel:empty_response" in str(e):
                    print(f"  ⚠ 检测到Gemini空响应错误，跳过重试...")
                    break  # 直接跳出重试循环
                
                # 检查是否是429错误（API负载过高）
                elif "429" in str(e) or "负载已饱和" in str(e) or "Too Many Requests" in str(e):
                    print(f"  ⚠ 检测到API负载过高，延长等待时间...")
                    delay = random.randint(120, 300)  # 2-5分钟
                else:
                    delay = random.randint(10, 20)  # 随机30-60秒
                
                # 如果不是最后一次尝试，等待后重试
                if attempt < max_retries - 1:
                    print(f"  等待 {delay} 秒后重试...")
                    time.sleep(delay)
        
        # 如果所有重试都失败了
        if not success:
            # 计算实际尝试次数
            actual_attempts = max_retries
            if "received empty response from Gemini" in str(last_error) or "no meaningful content in candidates" in str(last_error) or "channel:empty_response" in str(last_error):
                # 如果是空响应错误，实际尝试次数是当前尝试次数
                actual_attempts = attempt + 1 if 'attempt' in locals() else 1
            
            print(f"  ✗ 处理失败: {os.path.basename(image_path)} (尝试了 {actual_attempts} 次)")
            
            # 将失败的原图移动到other文件夹
            try:
                move_to_other_folder(image_path)
            except Exception as move_error:
                print(f"  ⚠ 移动失败图片到other文件夹失败: {move_error}")
            
            stats["failed"] += 1
            stats["details"].append({
                "input": image_path,
                "output": None,
                "status": "failed",
                "attempts": actual_attempts,
                "error": last_error
            })
        
        # 在处理之间添加随机延迟，避免API限制
        if i < len(image_files):
            delay = random.randint(5, 15)  # 增加到10-20秒
            print(f"  等待 {delay} 秒后处理下一张...")
            time.sleep(delay)
    
    return stats

def print_summary(stats):
    """打印处理结果摘要"""
    print("\n" + "="*60)
    print("批量处理完成！")
    print("="*60)
    print(f"总图片数: {stats['success'] + stats['failed'] + stats['skipped']}")
    print(f"成功处理: {stats['success']}")
    print(f"处理失败: {stats['failed']}")
    print(f"跳过文件: {stats['skipped']}")
    
    if stats['failed'] > 0:
        print("\n失败的图片:")
        for detail in stats['details']:
            if detail['status'] == 'failed':
                print(f"  - {os.path.basename(detail['input'])}: {detail.get('error', '未知错误')}")
    
    if stats['success'] > 0:
        print("\n成功生成的图片:")
        for detail in stats['details']:
            if detail['status'] == 'success':
                print(f"  - {os.path.basename(detail['output'])}")

if __name__ == "__main__":
    # 配置参数
    boots_image = "src/boots.png"
    custom_prompt = "图中女生穿上牛仔裤和图二中靴子, 保持脸部不变, 镜头拉远保证女生全身和靴子均在图内"
    
    # 检查靴子图片是否存在
    if not os.path.exists(boots_image):
        print(f"错误: 靴子图片文件不存在: {boots_image}")
        exit(1)
    
    print("开始批量处理图片...")
    print(f"靴子图片: {boots_image}")
    print(f"提示词: {custom_prompt}")
    print("-" * 60)
    
    # 执行批量处理
    results = batch_process_images_with_boots(
        boots_image_path=boots_image,
        prompt=custom_prompt,
        max_retries=3,
        timeout_seconds=240
    )
    
    # 打印结果摘要
    print_summary(results)
